import torch
import torch.utils.data as Data
from torchvision import transforms
from torchvision.datasets import FashionMNIST
from model import LeNet


def test_data_process():  # 训练——验证下载、处理函数
    test_data = FashionMNIST(root='D:/Train/data',
                              train=False,
                              transform=transforms.Compose([transforms.Resize(size=28), transforms.ToTensor()]),
                              download=True)

    # 数据打包
    test_dataloader = Data.DataLoader(dataset=test_data,
                                   batch_size=1,
                                   shuffle=True,
                                   num_workers=2)
    return test_dataloader


def test_model_process(model, test_dataloader):
    # 检测设备
    device = "cuda" if torch.cuda.is_available() else 'cpu'
    # 模型放到设备中
    model = model.to(device)

    # 初始化参数
    test_corrects = 0.0
    test_num = 0

    # 只进行前向传播，不计算梯度，从而节省内存，加快运行速度
    with torch.no_grad():
        for test_data_x, test_data_y in test_dataloader:
            # 数据放到设备中
            test_data_x = test_data_x.to(device)
            # 标签防爆设备中
            test_data_y = test_data_y.to(device)
            # 设置模型为评估模式
            model.eval()
            # 前向传播过程，输入为测试数据集，输出为对每个样本的预测值
            output = model(test_data_x)
            # 查找每一行最大值对应的行标
            pre_lab = torch.argmax(output, dim=1)
            # 若预测正确，则准确度test_corrects加1
            test_corrects += torch.sum(pre_lab == test_data_y.data)
            # 将所有的测试样本累加
            test_num += test_data_x.size(0)
        # 计算准确率
        test_acc = test_corrects.double().item() / test_num
        print("准确率：", test_acc)

if __name__ == "__main__":
    #加载模型
    model = LeNet()

    model.load_state_dict(torch.load('best_model.pth'))
    # 加载测试数据
    test_dataloader = test_data_process()
    #加载模型测试的函数
    test_model_process(model, test_dataloader)

    device = "cuda" if torch.cuda.is_available() else 'cpu'
    model = model.to(device)


    classes = ['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat', 'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot']
    with torch.no_grad():
        for b_x, b_y in test_dataloader:
            b_x = b_x.to(device)
            b_y = b_y.to(device)

            # 设置模型为验证模型
            model.eval()
            output = model(b_x)
            pre_lab = torch.argmax(output, dim=1)
            result = pre_lab.item()
            label = b_y.item()
            print("预测值为:", classes[result], "-------", "真实值：", classes[label])


